一、前言

繼上一篇將資料集建立完Dataframe後，接著要利用Scikit-learn將資料集拆分成訓練集、測試集以及驗證集，並建立成YOLO訓練的格式，那就開始吧。

二、拆分資料集

• 首先我們要先將資料集拆成訓練集以及測試集，這邊將資料集的10%當成測試的資料。

from sklearn.model_selection import train_test_split

# 將資料及拆成訓練以及測試集，並設定將資料集的10%作為測試的資料
train, test = train_test_split(alldata, test_size = 0.1, random_state=42)

# 再將資料進一步分成訓練以及驗證集
# 提取資料的8/9作為訓練用的資料，這樣整體的占比為 訓練/測試/驗證 = 80/10/10 
train, val = train_test_split(train, train_size = 8/9, random_state=42)

#輸出樣本的總數
print(f'''
      len(train) = {len(train)}
      len(val) = {len(val)}
      len(test) = {len(test)}
''')

輸出的結果為下圖，訓練集總共345筆，驗證及測試各佔44筆。

三、建立YOLO格式分類

• 接著將分割完的資料及進一步依照YOLO的格式進行分類。
• 分別建立3個資料夾，train, val, test。

# 設定資料夾的路徑以及新增資料夾
def make_split_folder_in_yolo_format(split_name, split_df):
    labels_path = os.path.join(dataset_path, 'Modle', split_name, 'labels')
    images_path = os.path.join(dataset_path, 'Modle', split_name, 'images')
   
    os.makedirs(labels_path, exist_ok=True)
    os.makedirs(images_path, exist_ok=True)

    # 對Datafram中的每一筆資料進行迭代
   for _, row in split_df.iterrows():
        img_name, img_extension = os.path.splitext(os.path.basename(row['img_path']))
        
        # 計算邊界框的數值
        x_center = (row['xmin'] + row['xmax']) / 2 / row['img_width']
        y_center = (row['ymin'] + row['ymax']) / 2 / row['img_height']
        width = (row['xmax'] - row['xmin']) / row['img_width']
        height = (row['ymax'] - row['ymin']) / row['img_height']

        # 將標籤儲存為YOLO的格式
        label_path = os.path.join(labels_path, f'{img_name}.txt')
        with open(label_path, 'w') as file:
            file.write(f"0 {x_center:.4f} {y_center:.4f} {width:.4f} {height:.4f}\n")
        
        #複製圖片到images的資料夾下
        shutil.copy(row['img_path'], os.path.join(images_path, img_name + img_extension))
    
    print(f"Created '{images_path}' and '{labels_path}'")

3. 接著執行make_split_folder_in_yolo_format副程式，建立YOLO格式的訓練集。

make_split_folder_in_yolo_format('train', train)
make_split_folder_in_yolo_format('val', val)
make_split_folder_in_yolo_format('test', test)

建立完成時如果有跳出下列訊息代表建立成功。

這樣我們就完成了樣本的分類，接著就可以進行YOLO的訓練了，那我們下篇見。